"""

Hierarchical clustering for ConsensusCluster


Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

import treetype
from scipy import cluster as cl


#def _assign_node_clust(node, id):
#    if node.value is not None:
#        node.cluster_id = id
#        return
#    else:
#        _assign_node_clust(node.left, id)
#        _assign_node_clust(node.right, id)
#
#        node.cluster_id = id

def _assign_tree_stack(node, id):
    # Dec 11 2009 - Using a recursive assignment hits the python recursion limit for tree depths
    #               greater than 1000, so we have to enumerate all the nodes in a stack

    remaining = [node]
    allnodes  = [node]

    while remaining:
        nd = remaining.pop()
        if nd.value is None:
            remaining.extend([nd.left, nd.right])
            allnodes.extend([nd.left, nd.right])

    for nd in allnodes:
        nd.cluster_id = id


def hierarchy(sdata, num_clusters, data_matrix=None, distance_matrix=False, distance_metric='euclidean', linkage='average', **kwds):
    """

    Wrapper for scipy hierarchical clustering and ConsensusCluster's treetype and display.Clustmap modules

    Note that the latter modules are deprecated and will likely be removed in favour of scipy's internal
    dendrogramming methods.

    """

    Tree = treetype.Tree

    if data_matrix is None:
        M = sdata.M
    else:
        M = data_matrix

    if not distance_matrix:
        Y = cl.hierarchy.distance.pdist(M, metric=distance_metric)
    else:
        Y = cl.hierarchy.distance.squareform(M, checks=False) 

    tree = [ Tree(value=x) for x in xrange(M.shape[0]) ]

    Z = cl.hierarchy.linkage(Y, method=linkage, metric=distance_metric)

    for i, j, dist, k in Z:
        tree.append(Tree(left=tree[int(i)], right=tree[int(j)], dist=dist))

    fclust = cl.hierarchy.fcluster(Z, num_clusters, criterion='maxclust')
    nodes, node_ids = cl.hierarchy.leaders(Z, fclust)

    #import sys
    #sys.setrecursionlimit(2000)

    for i in xrange(num_clusters):
        #_assign_node_clust(tree[nodes[i]], node_ids[i])
        _assign_tree_stack(tree[nodes[i]], node_ids[i])

    for i in xrange(M.shape[0]):
        sdata.samples[i].cluster_id = fclust[i]

    return tree[-1]
